#-*- coding: utf-8 -*-
"""
usage: python u7a3.py <im1> <im2> <im3>

Es werden 5 OBJ-Files im aktuellen Verzeichnis gespeichert.
Die 4 OBJs der 4-Scan Methode und das Ergebnis-OBJ der Oberflaeche.

@author: Pascal Trebert
@version: 1.0

2013 WS HSRM - Computer Vision Dr. Schwanecke - Blatt 07 A3

Shape From X Klasse 1:
Schritt 1: Obeflaechen-Normalen-Feld erzeugen -> Photometric Stereo
Schritt 2: Obeflaechen-Normalen-Feld integrieren -> 4-Scan-Methode (lokale Integration)

getestet auf Ubuntu 12.04 LTS, OBJs mit meshlab betrachtet.
Bilder: ellipsoid_-02_0_1.pgm ellipsoid_02_02_1.pgm ellipsoid_02_-02_1.pgm

"""
__author__ = 'p_trebert'

import sys
import warnings

import cv2
import numpy as np
import numpy.linalg as la

class FourScanMethod(object):
    """Integriert ein Normalenfeld mittels 4-Scan-Methode - Lokale 2D integration"""

    def __init__(self, normals_field):
        self.normals_field = normals_field
        self.gradients_fields = self.get_gradients_fields(normals_field)
        self.gradients_x_negative = self.get_gradients_fields(normals_field)[0]
        self.gradients_y_negative = self.get_gradients_fields(normals_field)[1]
        self.gradients_x = self.gradients_x_negative * -1
        self.gradients_y = self.gradients_y_negative * -1        # wichtig fuer Formel benutzen
        self.show('-p', self.gradients_x_negative)
        self.show('-q', self.gradients_y_negative)
        self.show('p', self.gradients_x)
        self.show('q', self.gradients_y)
        # Oberflaechen z1, z2, z3, z4
        self.surface_z1 = self.get_surface(self.normals_field, 'TL', self.gradients_x_negative, self.gradients_y)
        ## FIXME: q, p und -p, -q OK?? Was ist korrekt??
        self.surface_z2 = self.get_surface(self.normals_field, 'BL', self.gradients_x, self.gradients_y)
        # (-p, q), (p, -q) OK, finde p, -q besser
        self.surface_z3 = self.get_surface(self.normals_field, 'BR', self.gradients_x, self.gradients_y_negative)
        # p,q und -p,-q OK??
        self.surface_z4 = self.get_surface(self.normals_field, 'TR', self.gradients_x, self.gradients_y)
        self.surface = (self.surface_z1 + self.surface_z2 + self.surface_z3 + self.surface_z4) / 4.
        self.write_obj(self.surface, './surface.obj', self.normals_field)

    def write_obj(self, scalar_field, path, normals=None):
        """Berrechnet einen Obeflächenpunkt anhand lokaler Integrations-Methode (4 Scan)"""
        with file(path, 'w') as f:
            for index, val in np.ndenumerate(scalar_field):
                x, y = index
                f.write('v {} {} {}\n'.format(x, y, val))
            if normals is not None:
                assert(len(normals) == len(scalar_field))
                for index, val in np.ndenumerate(normals):
                    x, y, z = index
                    if z == 0:
                        f.write('vn {}'.format(val))
                    elif z == 1:
                        f.write(' {}'.format(val))
                    elif z == 2:
                        f.write(' {}\n'.format(val))
                    else:
                        raise IndexError("Ups, Normalen mal checken!")

    def get_iterators(self, width, height, start_pos):
        """Gibt korrekt laufende Iteratoren fuer x, y Achsen zurueck, je nach Startposition
        return -> langsam, schnell
        return -> iter_y, iter_x
        """
        if start_pos == 'TL':
            return xrange(1, height), xrange(1, width)
        if start_pos == 'BL':
            return xrange(1, width), xrange(height - 2, -1, -1)  # runterzaehlen
        if start_pos == 'BR':
            return xrange(width - 2, -1, -1), xrange(height - 2, -1, -1)  # runterzaehlen
        if start_pos == 'TR':
            return xrange(height - 2, -1, -1),  xrange(1, width)  # runterzaehlen

    def get_first_row_and_col(self, width, height, start_pos):
        """Gibt Indices der 'ersten' Zeile und Spalte zurueck anbhaengig von Startposition
        return -> (row1, line1)
        """
        if start_pos == 'TL':
            return 0, 0
        if start_pos == 'BL':
            return height - 1, 0
        if start_pos == 'BR':
            return height - 1, width - 1
        if start_pos == 'TR':
            return 0, width - 1

    def get_surface(self, normals_field, start_pos, gradients_x, gradients_y):
        """Gibt Hoehenfeld (Skalarfeld) zurueck, das anhand des Normalenfeld (Vektorfeld) und Startposition
        berrechnet wird.
        startpos: 'TL', 'TR', 'BL', 'BR'

        """
        #normals_field = np.floor(10 * np.random.random((4, 4, 3)))
        height, width = normals_field.shape[0], normals_field.shape[1]
        z = np.zeros((height, width), dtype=np.float64)
        fa, fb, fc = self.get_2x2_neighbour_positions(start_pos)
        iter_y, iter_x = self.get_iterators(width, height, start_pos)
        first_row, first_col = self.get_first_row_and_col(width, height, start_pos)
        z[0, 0] = 0                                              # Default-Startwert
        # TODO: erste reihe erste Spalte in extra schleifen machen
        for y in iter_y:
            for x in iter_x:
                langsam, schnell = y, x
                #print 'schnell %s, langsam: %s' %(schnell, langsam)
                a, b, c = fa(langsam, schnell), fb(langsam, schnell), fc(langsam, schnell)           # Bildpositionen der 2x2 Nachbarschaft als 2-Tuple
                #print 'a: %s, b: %s, c:%s' % (a, b, c)
                za = z[a]
                zc = z[c]

                if start_pos == 'BL' or start_pos == 'TR':  # Coloumn-Major
                    #print 'BL'
                    pxy, qxy = gradients_x[schnell, langsam], gradients_y[schnell, langsam]  # Bildgradienten
                    pc, qa = gradients_x[c], gradients_y[a]
                else:  # Row-Major UL, BR
                    pxy, qxy = gradients_x[langsam, schnell], gradients_y[langsam, schnell]  # Bildgradienten
                    pc, qa = gradients_x[c], gradients_y[a]
                if y == first_row:                                       # FIXME geht hir nicht rein row-major Erste Zeile -> Sonderrechnung
                    #print 'y ist erste reihe ', y
                    val = za + ((qa + qxy) / 2.)
                elif x == first_col:                                     # Erste Spalte -> Sonderrechnung
                    val = zc + ((pc + pxy) / 2.)
                else:                                            # Ab 2. zeile, 2. Spalte
                    pa, pb, = gradients_x[a], gradients_x[b]
                    qb, qc = gradients_y[b], gradients_y[c]
                    delta_x = (pb + pa + pc + pxy) / 4.
                    delta_y = (qb + qa + qc + qxy) / 4.
                    val = (za + zc + delta_x + delta_y) / 2
                if start_pos == 'BL' or start_pos == 'TR':
                    z[schnell, langsam] = val
                else:
                    z[langsam, schnell] = val

        self.write_obj(z, './{}.obj'.format(start_pos))                            # OBJ File schreiben
        return z

    def get_2x2_neighbour_positions(self, start_pos):
        """Definiert Funktionen fuer Nachbarn-Pixel-Positonen a, b, c der 2x2 Nachbarschaft, abhaengig von Startpunkt.
        return -> a(1,1), b(1,1), c(1,1) -> ((1, 0), (0, 0), (0, 1))
        numpy-Style enthalten die Tuple (Zeile, Spalte) -> a(y, x), b(y, x), c(y, x)

        """
        # hier kommt immer (langsam, schnell) rein
        if start_pos == 'TL':
            # Achtung: numpy-style y links, x rechts
            # Achtung: (langsam, schnell)
            a = lambda y, x: (y, x-1)
            b = lambda y, x: (y-1, x-1)
            c = lambda y, x: (y-1, x)
        elif start_pos == 'BL':
            # Achtung: (langsam, x schnell) rein
            a = lambda x, y: (y+1, x)
            b = lambda x, y: (y+1, x-1)
            c = lambda x, y: (y, x-1)
        elif start_pos == 'BR':
            # Achtung: (langsam, x schnell) rein
            a = lambda y, x: (y, x+1)
            b = lambda y, x: (y+1, x+1)
            c = lambda y, x: (y+1, x)
        elif start_pos == 'TR':
            # Achtung: (langsam, x schnell) rein
            a = lambda x, y: (y-1, x)
            b = lambda x, y: (y-1, x+1)
            c = lambda x, y: (y, x+1)
        return a, b, c

    def get_gradients_fields(self, normals_field):
        """Berrechnet 2 Gradientenfelder (Skalarfelder) in x oder y Richtung anhand des Normalenfelds.
        Durch Nutzung der Obeflaechennormalen und Kreuzprodukt sind die Bildgradienten p,q noch negativ.
        return -> -p, -q

        """
        n_xy = normals_field.take([0, 1], 2)                           # Array mit x,y Komponenten der Normalen
        n_z = normals_field.take([2], 2)                               # Array mit z Komponenten der Normalen
        with warnings.catch_warnings():                                # teilen durch 0 und teilen der 0 OK
            warnings.simplefilter("ignore")                            # -> warning unterdruecken
            pq = np.divide(n_xy, n_z)                                  # Gradienten -p = nx/nz, -q = ny/nz
        pq = np.where(np.isnan(pq), 0, pq)                             # ersetze nan und inf durch 0
        pq = np.where(np.isinf(pq), 0, pq)
        height, width = normals_field.shape[0], normals_field.shape[1]
        p = pq.take([0], 2).ravel(order='C').reshape((height, width))  # Gradientenfelder (Skalarfelder) p und q
        q = pq.take([1], 2).ravel(order='C').reshape((height, width))  # ravel->flat, reshape 2D
        return p, q

    def show(self, windowname, img):
        """Zeigt img in einem neuen Fenster windowname an"""
        cv2.namedWindow(windowname, cv2.CV_WINDOW_AUTOSIZE)
        cv2.imshow(windowname, img)


class PhotometricStereo():
    """Erzeugt das Obeflaechen-Normalen-Feld mittels Photometric Stereo - Shape From X Klasse 1: Schritt 1 von 2"""

    def __init__(self, pic_1, pic_2, pic_3):
        """3 Lichtquellen, Bildname muss Lichtposition enthalten z.B.'ellipsoid_02_02_1.pgm'"""
        # Initialisierungen der Bilddaten
        self.img_1 = self.read_image(pic_1)
        self.light_pos_1 = self.get_light_position(pic_1)
        self.img_1_windowname = 'Bild 1 %s' % self.light_pos_1
        self.show(self.img_1_windowname, self.img_1)

        self.img_2 = self.read_image(pic_2)
        self.light_pos_2 = self.get_light_position(pic_2)
        self.img_2_windowname = 'Bild 2 %s' % self.light_pos_2
        self.show(self.img_2_windowname, self.img_2)

        self.img_3 = self.read_image(pic_3)
        self.light_pos_3 = self.get_light_position(pic_3)
        self.img_3_windowname = 'Bild 3 %s' % self.light_pos_2
        self.show(self.img_3_windowname, self.img_3)

        # Initialisierung Photometric-Stereo Variablen
        self.lights_mat = self.get_lights_mat(self.light_pos_1, self.light_pos_2, self.light_pos_3)
        assert(la.det(self.lights_mat) != 0)
        self.lights_mat_inverted = self.lights_mat.I

        # Normalenfeld
        self.normal_field = self.get_normal_field()
        self.show('Normalen-Feld', self.normal_field)
        #print self.normal_field[50, 50]

    def get_normal_field(self, images=None, lights_mat=None):
        """Gibt das Normalenfeld zurueck, berechnet anhand Photometric Stereo"""
        #print 'get_normal_field()'
        images = [self.img_1, self.img_2, self.img_3] if images is None else images
        lights_mat = self.lights_mat if lights_mat is None else lights_mat
        assert(False not in [images[0].shape == i.shape for i in images[1:]])  # Test ob alles shapes gleich sind

        normal_field = np.ndarray(images[0].shape + (3,))  # 3D Array, (100, 100, 3), 3=>Normalenvektor
        #print 'SHAPE ', normal_field.shape
        height, width = images[0].shape
        for row in range(height):
            for col in range(width):
                normal = self.get_normal_normalized(col, row)  # muss normalisiert sein
                normal_field[row, col] = normal
        # test
        #self.compare_normals_at(50, 51, normal_field)
        return normal_field

    def get_lights_mat(self, s1, s2, s3):
        """Gibt Matrix mit Beleuchtungseigenschaften s1, s2, s3 zurueck"""
        if False in map(isinstance, [s1, s2, s3], [np.ndarray, np.ndarray, np.ndarray]):
            raise ValueError("%s->%s %s->%s %s->%s müssen numpy-Arrays sein!",
                             (s1, type(s1), s2, type(s2), s3, type(s3)))
        s_mat = np.concatenate((s1, s2, s3), axis=0)  # Achse 0 => Y , C-Style row major
        s_mat = s_mat.reshape(3, -1)  # reshape, da nur eindimensionale Arrays verknüpft wurden
        s_mat = np.asmatrix(s_mat, dtype=float)
        return s_mat

    def get_intensities(self, x, y, images=None):
        """Gibt Vektor mit allen Intensitaeten einer Bildposition zurueck"""
        assert(x < self.img_1.shape[1] and y < self.img_1.shape[0])
        images = [self.img_1, self.img_2, self.img_3] if images is None else images

        intensities = []
        for i in images:
            intensities.append(i[y, x])  # numpy -> rows, cols TODO:STIMMT DAS?
        return np.array(intensities, dtype=np.uint8)

    def get_normal(self, x, y, images=None, lights_mat=None):
        """Gibt den NICHT normalisierten Normalenvektor, berrechnet anhand aller Quellbilder,
        zu einer Bildposition zurueck

        """
        images = [self.img_1, self.img_2, self.img_3] if images is None else images
        lights_mat = self.lights_mat if lights_mat is None else lights_mat

        intensities = self.get_intensities(x, y, images)
        lights_mat_inverted = lights_mat.I

        normal = lights_mat_inverted.dot(intensities)
        normal = np.squeeze(np.asarray(normal))  # mat -> array, squeezed
        # normal
        return normal

    def get_albedo(self, x, y, images=None, lights_mat=None):
        """Gibt die Laenge des Normalen-Vektors (Rueckstrahleigenschaft) zu einer Bildposition zurueck"""
        images = [self.img_1, self.img_2, self.img_3] if images is None else images
        lights_mat = self.lights_mat if lights_mat is None else lights_mat

        normal = self.get_normal(x, y, images, lights_mat)
        albedo = la.norm(normal)  # TODO: Rundungsfehler??
        #print 'albedo ', albedo
        return albedo

    def get_normal_normalized(self, x, y, images=None, lights_mat=None):
        """Gibt den normalisierten Normalenvektor, berrechnet anhand aller Quellbilder, zu einer Bildposition zurück"""
        images = [self.img_1, self.img_2, self.img_3] if images is None else images
        lights_mat = self.lights_mat if lights_mat is None else lights_mat

        normal = self.get_normal(x, y, images, lights_mat)
        albedo = self.get_albedo(x, y, images, lights_mat)  # rueckstrahleigenschaft == Laenge der Normalen
        normal_normalized = normal / albedo if albedo != 0 else normal * 0  # div 0 -> NAN
        return normal_normalized

    def compare_normals_at(self, r, c, normal_field):
        """Vergleicht gesetzte Normalen im Normalenfeld mit direkter Berrechnung"""
        print
        print 'normalfield row:[%s] col:[%s] = %s' % (r, c, normal_field[r, c])
        print 'direkt berr.row:%s, col:%s = %s' % (r, c, self.get_normal(c, r))
        print 'Vertauschte Indizes:'
        print 'normalfield.: row:%s, col:%s: %s' % (c, r, normal_field[c, r])
        print 'direkt berr.: row:%s, col:%s: %s' % (c, r, self.get_normal(r, c))

    # Hilfsmethoden
    # TODO: Entfernen in Abgabe
    def show_gradients(self, normal_field, x=45, y=45):
        """Gibt auf Konsole die Gradienten -p, -q der Bildposition x, y aus.
        Erst anhand normalisierter Normalen dann anhand nicht normalisierter Normalen (ähnlich im Skript)

        """
        print 'normalized Normalfield[%s, %s] = %s' % (x, y, normal_field[x, y])
        print 'nnx = %s' % (normal_field[x, y][0])
        print 'nny = %s' % (normal_field[x, y][1])
        print 'nnz = %s' % (normal_field[x, y][2])
        print 'nx/nz = -p = %s' % (normal_field[x, y][0]/normal_field[x, y][2])
        print 'ny/nz = -q = %s' % (normal_field[x, y][1]/normal_field[x, y][2])
        print
        n = self.get_normal(x, y, [self.img_1, self.img_2, self.img_3])
        print 'Normalfield[%s, %s] = %s' % (x, y, n)
        print 'nx = %s' % (n[0])
        print 'ny = %s' % (n[1])
        print 'nz = %s' % (n[2])
        print 'nx/nz = -p = %s' % (n[0]/n[2])
        print 'ny/nz = -q = %s' % (n[1]/n[2])
        print

    def color_2_gray(self, img):
        return cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)

    # TODO: Entfernen in Abgabe
    def print_mat_infos(self, mat, name=None):
        """Zeigt Informationen zu einer numpy-Matrix"""
        assert(isinstance(mat, np.matrix))
        print '\n-----------------------------------'
        print 'Matrix-Bezeichnung: ', name
        print 'matrix: \n', mat
        print '\ndeterminante: ', la.det(mat)
        print '\ninvertierte matrix: \n', mat.I
        print '\ntype matrix=%s\btype invertierte matrix=%s' % (type(mat), type(mat.I))
        print '\nmat * matinv\n', mat * mat.I
        print '-----------------------------------\n'

    def read_image(self, file, method=cv2.CV_LOAD_IMAGE_COLOR):
        """Liest eine Bilddatei von Platte in gewuenschtem Farbmodus"""
        img = cv2.imread(file, method)
        gray = self.color_2_gray(img)
        return gray

    def show(self, windowname, img):
        """Zeigt img in einem neuen Fenster windowname an"""
        cv2.namedWindow(windowname, cv2.CV_WINDOW_AUTOSIZE)
        cv2.imshow(windowname, img)

    def get_light_position(self, string):
        """"ellipsoid_02_02_1.pgm -> array([2, 2, 1])"""
        a, b, c = string[:-4].split('_')[-3:]
        pos = []
        for i in a, b, c:
            i_s = i.strip('-')
            vz = -1 if i[0] == '-' else 1
            if len(i_s) > 1 and i_s[0] == '0':
                pos.append(vz * float(i_s[1]) * 10**-1)
            else:
                pos.append(vz * float(i_s))
        return np.array(pos)

    def close(self):
        """Schliesst alle Fenster"""
        for w in self.img_1_windowname, self.img_2_windowname, self.img_3_windowname:
            cv2.destroyWindow(w)


def main(argv):
    if len(argv) == 0:
        # TODO: usage -> rtfm
        pic_1 = '../img/ellipsoid_-02_0_1.pgm'
        pic_2 = '../img/ellipsoid_02_02_1.pgm'
        pic_3 = '../img/ellipsoid_02_-02_1.pgm'
        argv = [pic_1, pic_2, pic_3]

    ps = PhotometricStereo(*argv)
    #ps.show_gradients()
    fsm = FourScanMethod(ps.normal_field)


    cv2.waitKey(0)  # unendlich warten auf Tastendruck
    ps.close()
    cv2.destroyAllWindows()
    sys.exit(0)

if __name__ == '__main__':
    main(sys.argv[1:])